# -*- coding: utf-8 -*-
"""
@Time: 2021/1/12 下午 8:16
@Author: jinzhuan
@File: interface.py
@Desc: 
"""
import torch
from cognlp import *
from cognlp.utils.util import load_yaml
from flask import Flask, jsonify, request
from flask_cors import *

torch.cuda.set_device(0)
config = load_yaml('../config/interface.yaml')
device = torch.device(config['device'])

tokenize_toolkit = TokenizeToolkit()

ner_toolkit = NerToolkit(
    bert_model=config['bert_model'],
    model_path=config['ner_model_path'],
    vocabulary_path=config['ner_vocabulary_path'],
    device=device,
    device_ids=config['device_id'],
    max_seq_length=config['max_seq_length'])

et_toolkit = EtToolkit(
    bert_model=config['bert_model'],
    model_path=config['et_model_path'],
    vocabulary_path=config['et_vocabulary_path'],
    device=device,
    device_ids=config['device_id'],
    max_seq_length=config['max_seq_length'])

el_toolkit = ElToolkit(wikidata2wikipedia_path=config['wikidata2wikipedia_path'])

re_toolkit = ReToolkit(
    bert_model=config['bert_model'],
    model_path=config['re_model_path'],
    vocabulary_path=config['re_vocabulary_path'],
    device=device,
    device_ids=config['device_id'],
    max_seq_length=config['max_seq_length'])

fn_toolkit = FnToolkit(
    bert_model=config['bert_model'],
    model_path=config['fn_model_path'],
    vocabulary_path=config['fn_vocabulary_path'],
    device=device,
    device_ids=config['device_id'],
    max_seq_length=config['max_seq_length'])

ee_ner_toolkit = NerAce2005Toolkit(
    bert_model=config['bert_model'],
    model_path=config['ee_ner_model_path'],
    vocabulary_path=config['ee_ner_vocabulary_path'],
    device=device,
    device_ids=config['device_id'],
    max_seq_length=config['max_seq_length'])

ee_toolkit = EeToolkit(
    bert_model=config['bert_model'],
    model_path=config['ee_model_path'],
    trigger_vocabulary_path=config['ee_trigger_vocabulary_path'],
    argument_vocabulary_path=config['ee_argument_vocabulary_path'],
    device=device,
    device_ids=config['device_id'],
    max_seq_length=config['max_seq_length'])

argument_toolkit = ArgumentToolkit(
    bert_model=config['bert_model'],
    model_path=config['argument_model_path'],
    device=device,
    device_ids=[4],
    max_seq_length=config['device_id'],
    trigger_vocabulary_path=config['frame_vocabulary_path'],
    argument_vocabulary_path=config['argument_vocabulary_path'])

app = Flask(__name__)
CORS(app, supports_credentials=True)
app.config['JSON_AS_ASCII'] = False


@app.route('/ner', methods=["GET", "POST"])
def ner():
    sentence = request.values.get('sentence')
    words = tokenize_toolkit.run(sentence)
    ner_result = ner_toolkit.run(words)
    return jsonify({"words": words, "ner_result": ner_result})


@app.route('/typing', methods=["GET", "POST"])
def typing():
    sentence = request.values.get('sentence')
    words = tokenize_toolkit.run(sentence)
    ner_result = ner_toolkit.run(words)
    et_result = et_toolkit.run(words, ner_result)
    return jsonify({"words": words, "ner_result": ner_result, "et_result": et_result})


@app.route('/linking', methods=["GET", "POST"])
def linking():
    sentence = request.values.get('sentence')
    words = tokenize_toolkit.run(sentence)
    el_result = el_toolkit.run(sentence)
    return jsonify({"words": words, "el_result": el_result})


@app.route('/relation', methods=["GET", "POST"])
def relation():
    sentence = request.values.get('sentence')
    words = tokenize_toolkit.run(sentence)
    ner_result = ner_toolkit.run(words)
    re_result = re_toolkit.run(words, ner_result)
    return jsonify({"words": words, "ner_result": ner_result, "re_result": re_result})


@app.route('/frame', methods=["GET", "POST"])
def frame():
    sentence = request.values.get('sentence')
    words = tokenize_toolkit.run(sentence)
    fn_result = fn_toolkit.run(words)
    element_result = argument_toolkit.run(words, fn_result)
    return jsonify({"words": words, "fn_result": fn_result, 'element_result': element_result})


@app.route('/event', methods=["GET", "POST"])
def event():
    sentence = request.values.get('sentence')
    words = tokenize_toolkit.run(sentence)
    ner_result = ee_ner_toolkit.run(words)
    ee_result = ee_toolkit.run(words, ner_result)
    return jsonify({"words": words, "ner_result": ner_result, "ee_result": ee_result})


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=9988)
